import numpy as np
import matplotlib.pyplot as plt
import math
from scipy import integrate
import mpl_toolkits.axisartist as axisartist


def func1(t, A, f, fi, mode):
    # 振幅为A、频率为f, 初相为fi的正弦/余弦信号，mode=sin/cos
    fi = fi * math.pi / 180
    fi_prime = 0 if mode == 'sin' else math.pi/2
    return A * math.sin(2 * math.pi * f * t + fi + fi_prime)

def func2(t):
    # 分段函数f(t)
    if t <= 1:
        return 0
    elif t <= 2:
        return t - 1
    elif t <= 3:
        return 1
    elif t <= 4:
        return t - 2
    else:
        return 0

def si(t):
    def sa(x):
        return math.sin(x) / x
    return integrate.quad(sa, 0, t)
    

def set_axis(fig, number, title, range=(-12, 12)):
    ax = axisartist.Subplot(fig, number)
    fig.add_axes(ax)
    ax.axis[:].set_visible(False)
    ax.axis["x"] = ax.new_floating_axis(0,0)
    ax.axis["x"].set_axisline_style("->", size = 1.0)
    ax.axis["y"] = ax.new_floating_axis(1,0)
    ax.axis["y"].set_axisline_style("->", size = 1.0)
    ax.axis["x"].set_axis_direction("bottom")
    ax.axis["y"].set_axis_direction("left")
    ax.set_xlim(range[0], range[1])
    ax.set_title(title)


if __name__ == "__main__":

    t = np.linspace(-10, 10, 100, True, False, float, 0)

    fig = plt.figure(figsize=(10, 7))
    plt.rcParams['font.family'] = ['Arial Unicode MS']

    # 绘制振幅为2，频率为10Hz的正弦信号
    f1 = np.array([func1(i, 2, 10, 0, 'sin') for i in t])
    set_axis(fig, 321, "振幅为2、频率为10Hz的正弦信号")
    plt.plot(t, f1)

    # 绘制振幅为4，频率为20Hz，初相为50度的余弦信号
    f2 = np.array([func1(i, 4, 20, 50, 'cos') for i in t])
    set_axis(fig, 322, "振幅为4、频率为20Hz、初相为50的余弦信号")
    plt.plot(t, f2)
    
    # 绘制上述信号相加的结果
    f3 = f1 + f2
    set_axis(fig, 323, "信号相加")
    plt.plot(t, f3)

    # 绘制连续信号si(t)
    f4 = np.array([si(i)[0] for i in t])
    set_axis(fig, 324, "Si(t)信号")
    plt.plot(t, f4)

    t = np.linspace(-10, 10, 1000, True, False, float, 0)
    
    # 绘制原始分段函数f(t)
    f5 = np.array([func2(i) for i in t])
    set_axis(fig, 325, "f(t)图像", (0, 5))
    plt.plot(t, f5)

    # 绘制变换后的分段函数f(3-2t)
    f6 = np.array([func2(3 - 2 * i) for i in t])
    set_axis(fig, 326, "f(3-2t)图像", (-3, 2))
    plt.plot(t, f6)

    plt.show()